use crate::util::{property, rerank};
use colored::Colorize;
use itertools::Itertools;
use rig::client::EmbeddingsClient;
use rig::vector_store::VectorSearchRequest;
use rig::{completion::ToolDefinition, tool::Tool};
use rig::{
    providers::{self, ollama::Client},
    vector_store::VectorStoreIndex,
    vector_store::VectorStoreIndexDyn,
    Embed, OneOrMany,
};
use rig_sqlite::{
    Column, ColumnValue, SqliteVectorIndex, SqliteVectorStore, SqliteVectorStoreTable,
};
use rusqlite::ffi::sqlite3_auto_extension;
use serde::{Deserialize, Serialize};
use serde_json::json;
use sqlite_vec::sqlite3_vec_init;
use tokio::task;
use tokio_rusqlite::Connection;

#[derive(Deserialize)]
pub struct OperationArgs {
    search_item: String,
}

#[derive(Debug, thiserror::Error)]
#[error("Math error")]
pub struct CatError;

#[derive(Deserialize, Serialize)]
pub struct SearchCmd {
    pub number: usize,
    pub vector_db_path: String,
}

impl Tool for SearchCmd {
    const NAME: &'static str = "rag_query";

    type Error = CatError;
    type Args = OperationArgs;
    type Output = String;

    async fn definition(&self, _prompt: String) -> ToolDefinition {
        ToolDefinition {
            name: "rag_query".to_string(),
            description: format!("搜索知识库，每次返回{}条", self.number),
            parameters: json!({
                "type": "object",
                "properties": {
                    "search_item": {
                        "type": "string",
                        "description": "需要搜索的内容"
                    }
                }
            }),
        }
    }

    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
        // todo 太复杂了
        let search_item = args.search_item.clone();
        let number = self.number;
        let vector_db_path = self.vector_db_path.clone();
        if let Ok(results) =
            task::spawn_blocking(move || search(&search_item, number, &vector_db_path)).await
        {
            // 合并结果
            let result_str = results
                .iter()
                .map(|item: &(f64, String, serde_json::Value)| {
                    let v: &serde_json::Value = &item.2;
                    format!(
                        "\n{}\n{}\n```\n{}\n```",
                        v["id"].as_str().unwrap_or("").red(),
                        v["desc"].as_str().unwrap_or(""),
                        v["content"].as_str().unwrap_or("").green()
                    )
                })
                .join("\n");

            Ok(result_str)
        } else {
            Ok("tool call fail".to_string())
        }
    }
}

// todo 冗余
#[derive(Embed, Clone, Debug, Deserialize)]
struct Document {
    id: String,
    content: String,
    #[embed]
    desc: String,
}

impl SqliteVectorStoreTable for Document {
    fn name() -> &'static str {
        "documents"
    }

    fn schema() -> Vec<Column> {
        vec![
            Column::new("id", "TEXT PRIMARY KEY"),
            Column::new("content", "TEXT"),
            Column::new("desc", "TEXT"),
        ]
    }

    fn id(&self) -> String {
        self.id.clone()
    }

    fn column_values(&self) -> Vec<(&'static str, Box<dyn ColumnValue>)> {
        vec![
            ("id", Box::new(self.id.clone())),
            ("content", Box::new(self.content.clone())),
            ("desc", Box::new(self.desc.clone())),
        ]
    }
}

fn search(
    search_item: &String,
    number: usize,
    vector_db_path: &String,
) -> Vec<(f64, String, serde_json::Value)> {
    unsafe {
        sqlite3_auto_extension(Some(std::mem::transmute(sqlite3_vec_init as *const ())));
    }

    // 阻塞运行
    let rt = tokio::runtime::Runtime::new().unwrap();

    // 大模型运行
    let env: property::Env = property::Env::init();
    let client = Client::builder().base_url(&env.embedding_endpoint).build();
    let model = client.embedding_model_with_ndims(&env.embedding_model, 1024);
    let conn = rt
        .block_on(Connection::open(vector_db_path))
        .expect("get connect fail");

    // 打开存储库
    let vector_store: SqliteVectorStore<
        providers::ollama::EmbeddingModel<reqwest::Client>,
        Document,
    > = rt
        .block_on(SqliteVectorStore::new(conn, &model))
        .expect("open sqlite db fail");

    // 搜索
    let query_req = VectorSearchRequest::builder()
        .samples(number as u64 * 10)
        .query(search_item)
        .build()
        .unwrap();
    let index = vector_store.index(model);
    let results = rt
        .block_on(index.top_n(query_req))
        .expect("get search resul fail");

    // reranker重排序
    let desc_list: Vec<String> = results
        .iter()
        .map(|doc| {
            let v: &serde_json::Value = &doc.2;
            v["desc"].as_str().unwrap_or("").to_string().clone()
        })
        .collect();
    let reranker = rerank::Reranker {
        end_point: env.rerank_endpoint.clone(),
    };
    let results: Vec<(f64, String, serde_json::Value)> =
        rt.block_on(reranker.rerank(search_item.to_string(), desc_list.clone(), results.clone()));
    let mut results: Vec<(f64, String, serde_json::Value)> =
        results.into_iter().take(number).collect();

    results
}

#[tokio::test]
async fn test_search() {
    let search_cmd = SearchCmd {
        number: 3,
        vector_db_path: "vector.db".to_string(),
    };

    let result_str = search_cmd
        .call(OperationArgs {
            search_item: "权限校验".to_string(),
        })
        .await
        .expect("执行失败");
    println!("result_str {}", result_str);
}
